#include <chrono>
#include <fstream>
#include <iostream>
#include <opencv2/opencv.hpp>
#include <string>
#include <thread>
#include <vector>

#include "common.hpp"
#include "det/postprocess.h"
#include "det/preprocess.h"
#include "det/rkpt.hpp"

// 配置参数 - 与模型匹配的尺寸
#define MODEL_INPUT_WIDTH 640
#define MODEL_INPUT_HEIGHT 480
#define CAMERA_WIDTH 640
#define CAMERA_HEIGHT 480
#define CONF_THRESHOLD 0.25
#define NMS_THRESHOLD 0.45

class CameraDetector
{
  private:
    cv::VideoCapture cap;
    RkPt *rknn_model;
    int model_input_h;
    int model_input_w;
    float conf_threshold;
    float nms_threshold;

  public:
    CameraDetector(const std::string &model_path, int camera_id = 0)
        : model_input_h(MODEL_INPUT_HEIGHT), model_input_w(MODEL_INPUT_WIDTH), conf_threshold(CONF_THRESHOLD),
          nms_threshold(NMS_THRESHOLD)
    {

        // 初始化摄像头
        cap.open(camera_id);
        if (!cap.isOpened())
        {
            std::cerr << "无法打开摄像头 " << camera_id << std::endl;
            std::cerr << "caonima " << camera_id << std::endl;
            exit(-1);
        }

        // 设置摄像头参数
        cap.set(cv::CAP_PROP_FRAME_WIDTH, CAMERA_WIDTH);
        cap.set(cv::CAP_PROP_FRAME_HEIGHT, CAMERA_HEIGHT);
        cap.set(cv::CAP_PROP_FPS, 30);

        // 初始化RKNN模型
        rknn_model = new RkPt(model_path);
        if (rknn_model->init(nullptr, false) != 0)
        {
            std::cerr << "RKNN模型初始化失败" << std::endl;
            exit(-1);
        }

        std::cout << "摄像头检测器初始化成功" << std::endl;
    }

    ~CameraDetector()
    {
        if (cap.isOpened())
        {
            cap.release();
        }
        if (rknn_model)
        {
            delete rknn_model;
        }
    }

    void run()
    {
        cv::Mat frame;
        int frame_count = 0;
        auto start_time = std::chrono::high_resolution_clock::now();

        std::cout << "开始实时检测，按 'q' 退出..." << std::endl;

        while (true)
        {
            try
            {
                cap >> frame;
                if (frame.empty())
                {
                    std::cerr << "无法获取摄像头帧" << std::endl;
                    std::this_thread::sleep_for(std::chrono::milliseconds(100));
                    continue;
                }

                // 检测当前帧
                DetectionResultsGroup results;
                if (detect_frame(frame, results))
                {
                    // 显示结果
                    show_results(frame, results, frame_count);
                }

                // 显示原始帧
                cv::imshow("Material Detection", frame);

                // 计算FPS
                frame_count++;
                if (frame_count % 30 == 0)
                {
                    auto current_time = std::chrono::high_resolution_clock::now();
                    auto duration = std::chrono::duration_cast<std::chrono::milliseconds>(current_time - start_time);
                    float fps = 30000.0f / duration.count();
                    std::cout << "FPS: " << fps << std::endl;
                    start_time = current_time;
                }

                // 检查退出条件
                char key = cv::waitKey(1);
                if (key == 'q' || key == 27)
                { // 'q' 或 ESC
                    break;
                }
            }
            catch (const std::exception &e)
            {
                std::cerr << "主循环中发生异常: " << e.what() << std::endl;
                std::this_thread::sleep_for(std::chrono::milliseconds(100));
                continue;
            }
            catch (...)
            {
                std::cerr << "主循环中发生未知异常" << std::endl;
                std::this_thread::sleep_for(std::chrono::milliseconds(100));
                continue;
            }
        }
    }

  private:
    bool detect_frame(cv::Mat &frame, DetectionResultsGroup &results)
    {
        try
        {
            // 检查输入帧是否有效
            if (frame.empty())
            {
                std::cerr << "输入帧为空" << std::endl;
                return false;
            }

            // 检查RKNN模型是否有效
            if (rknn_model == nullptr)
            {
                std::cerr << "RKNN模型未初始化" << std::endl;
                return false;
            }

            // 直接使用原始帧进行推理，不进行预处理
            // letterbox处理已经在rknn_model->infer内部实现
            results = rknn_model->infer(frame, 0);

            // 设置当前图像和帧ID
            results.cur_img = frame.clone();
            results.cur_frame_id = 0;

            return true;
        }
        catch (const std::exception &e)
        {
            std::cerr << "检测过程中发生异常: " << e.what() << std::endl;
            return false;
        }
        catch (...)
        {
            std::cerr << "检测过程中发生未知异常" << std::endl;
            return false;
        }
    }

    void show_results(cv::Mat &frame, const DetectionResultsGroup &results, int frame_id)
    {
        // 添加调试信息
        std::cout << "=== 显示检测结果 ===" << std::endl;
        std::cout << "检测到 " << results.dets.size() << " 个目标" << std::endl;

        // 在图像上绘制检测结果
        for (size_t i = 0; i < results.dets.size(); i++)
        {
            const auto &det = results.dets[i];
            
            // 基本有效性检查
            if (det.box.width <= 0 || det.box.height <= 0) {
                std::cout << "跳过无效检测框 " << i << ": " << det.box << std::endl;
                continue;
            }
            
            // 边界检查
            if (det.box.x < 0 || det.box.y < 0 || 
                det.box.x + det.box.width > frame.cols || 
                det.box.y + det.box.height > frame.rows) {
                std::cout << "检测框 " << i << " 超出图像边界，进行裁剪: " << det.box << std::endl;
                
                // 裁剪到图像边界内
                cv::Rect clipped_box;
                clipped_box.x = std::max(0, det.box.x);
                clipped_box.y = std::max(0, det.box.y);
                clipped_box.width = std::min(det.box.width, frame.cols - clipped_box.x);
                clipped_box.height = std::min(det.box.height, frame.rows - clipped_box.y);
                
                if (clipped_box.width <= 0 || clipped_box.height <= 0) {
                    std::cout << "裁剪后框无效，跳过" << std::endl;
                    continue;
                }
                
                std::cout << "绘制裁剪后的检测框: " << det.det_name << " 置信度: " << det.score 
                         << " 位置: " << clipped_box << std::endl;

                // 绘制边界框
                cv::rectangle(frame, clipped_box, cv::Scalar(0, 255, 0), 2);

                // 绘制标签和置信度
                std::string label = det.det_name + " " + std::to_string(det.score).substr(0, 4);
                cv::putText(frame, label, cv::Point(clipped_box.x, clipped_box.y - 10), cv::FONT_HERSHEY_SIMPLEX, 0.5,
                            cv::Scalar(0, 255, 0), 2);

                // 绘制目标中心点
                cv::Point center(clipped_box.x + clipped_box.width / 2, clipped_box.y + clipped_box.height / 2);
                cv::circle(frame, center, 5, cv::Scalar(0, 0, 255), -1);
            } else {
                std::cout << "绘制检测框: " << det.det_name << " 置信度: " << det.score 
                         << " 位置: " << det.box << std::endl;

                // 绘制边界框
                cv::rectangle(frame, det.box, cv::Scalar(0, 255, 0), 2);

                // 绘制标签和置信度
                std::string label = det.det_name + " " + std::to_string(det.score).substr(0, 4);
                cv::putText(frame, label, cv::Point(det.box.x, det.box.y - 10), cv::FONT_HERSHEY_SIMPLEX, 0.5,
                            cv::Scalar(0, 255, 0), 2);

                // 绘制目标中心点
                cv::Point center(det.box.x + det.box.width / 2, det.box.y + det.box.height / 2);
                cv::circle(frame, center, 5, cv::Scalar(0, 0, 255), -1);
            }
        }
        
        std::cout << "=== 检测结果显示完成 ===" << std::endl;
    }
};

int main(int argc, char *argv[])
{
    std::string model_path = "../model/wuzi_640x640.rknn"; // 默认模型路径
    int camera_id = 0;                             // 默认摄像头ID

    // 解析命令行参数
    if (argc > 1)
    {
        model_path = argv[1];
    }
    if (argc > 2)
    {
        camera_id = std::stoi(argv[2]);
    }

    // 检查模型文件是否存在
    std::ifstream model_file(model_path);
    if (!model_file.good())
    {
        std::cerr << "错误：模型文件不存在: " << model_path << std::endl;
        std::cerr << "请确保模型文件路径正确" << std::endl;
        return -1;
    }
    model_file.close();

    std::cout << "使用模型: " << model_path << std::endl;
    std::cout << "使用摄像头: " << camera_id << std::endl;

    try
    {
        CameraDetector detector(model_path, camera_id);
        detector.run();
    }
    catch (const std::exception &e)
    {
        std::cerr << "错误: " << e.what() << std::endl;
        return -1;
    }

    return 0;
}